from io import BytesIO
import socket
import struct
import os
import logging
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor
from PIL import Image

from common import db
from common import llm
from common import tomodel as tm
from conf import operationConfig
from conf import setting

model = tm.TomotaModel.load_model()
llm_client = llm.DeepSeekClient.get_instance()
database = db.Database.get_instance()

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    # filename='server.log'
)
# console = logging.StreamHandler()
# console.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
# console.setFormatter(formatter)
# logging.getLogger('').addHandler(console)

# 创建保存目录
os.makedirs(setting.FILE_PATH["SAVE_IMAGE"], exist_ok=True)


def recv_exact(sock, length):
    """确保接收指定长度的字节数据"""
    data = b''
    while len(data) < length:
        remaining = length - len(data)
        chunk = sock.recv(min(setting.BUFFER_SIZE, remaining))
        if not chunk:
            # 客户端断开连接
            raise ConnectionResetError("客户端连接断开")
        data += chunk
    return data


def handle_client(client_socket, client_address):
    """处理客户端连接"""
    logging.info(f"客户端已连接: {client_address}")

    try:
        with client_socket:
            # 读取数据包总长度
            total_length_data = recv_exact(client_socket, 4)
            total_length = struct.unpack('!I', total_length_data)[0]
            logging.debug(f"收到总长度: {total_length} 字节")

            # 读取参数1（float）
            param1_data = recv_exact(client_socket, 4)
            confidence = struct.unpack('!f', param1_data)[0]

            # 读取参数2（float）
            param2_data = recv_exact(client_socket, 4)
            iou = struct.unpack('!f', param2_data)[0]

            # 读取图片长度
            image_length_data = recv_exact(client_socket, 4)
            image_length = struct.unpack('!I', image_length_data)[0]
            logging.info(f"收到参数: param1={confidence:.2f}, param2={iou:.2f}, 图片大小={image_length} 字节")

            # 读取图片数据
            image_data = b''
            remaining = image_length
            while remaining > 0:
                chunk = client_socket.recv(min(setting.BUFFER_SIZE, remaining))
                if not chunk:
                    raise ConnectionResetError("客户端连接断开")
                image_data += chunk
                remaining -= len(chunk)

            logging.info(f"图片接收完成，实际大小: {len(image_data)} 字节")

            # # 保存图片
            file_path = save_image(image_data)
            logging.info(f"图片保存路径: {file_path}")

            # 处理数据
            processed_text, processed_image = process_data(image_data, confidence, iou)

            # 发送响应
            send_response(client_socket, processed_text, processed_image)
            logging.info(f"响应已发送给客户端: {client_address}")

    except ConnectionResetError as e:
        logging.error(f"客户端连接重置: {e}")
    except struct.error as e:
        logging.error(f"数据解析错误: {e}")
    except Exception as e:
        logging.error(f"处理客户端请求时出错: {e}", exc_info=True)
    finally:
        logging.info(f"客户端连接已关闭: {client_address}")


def process_data(image_data, param1, param2):
    """处理数据（示例：根据参数生成处理后的图片和文本）"""
    try:
        image = Image.open(BytesIO(image_data))
        
        # 如果模型要求输入为 RGB 格式，可强制转换
        if image.mode != 'RGB':
            image = image.convert('RGB')
        processed_image, processed_text = model.predict(image, param1, param2)
        # 将processed_image 转换为 bytes 类型
        processed_image_bytes = BytesIO().getvalue()
        
        processed_text = processed_text + "\n" + (llm_client.get_suggestion_detection_result(processed_text) or "")
        
        return processed_text, processed_image_bytes

    except Exception as e:
        logging.error(f"处理数据时出错: {e}", exc_info=True)
        return f"处理失败: {str(e)}", b''


def save_image(image_data):
    """保存图片到文件"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    file_path = os.path.join(setting.FILE_PATH["SAVE_IMAGE"], f"received_image_{timestamp}.png")

    try:
        with open(file_path, 'wb') as f:
            f.write(image_data)
        return file_path
    except Exception as e:
        logging.error(f"保存图片失败: {e}", exc_info=True)
        return "保存失败"


def send_response(client_socket, processed_text, processed_image):
    """发送响应给客户端"""
    try:
        # 发送处理后的文本
        text_bytes = processed_text.encode('utf-8')
        client_socket.sendall(struct.pack('!I', len(text_bytes)))
        client_socket.sendall(text_bytes)

        # 发送处理后的图片
        client_socket.sendall(struct.pack('!I', len(processed_image)))
        client_socket.sendall(processed_image)

    except Exception as e:
        logging.error(f"发送响应失败: {e}", exc_info=True)
        raise


def main():
    conf = operationConfig.OperationConfig()
    """主函数：启动服务器"""
    bindip = conf.get_section_for_data("SERVER", "bind")
    port = int(conf.get_section_for_data("SERVER", "port"))
    logging.info(f"服务器启动中，监听IP: {bindip}，端口: {port}")
    logging.info(f"图片保存目录: {setting.FILE_PATH["SAVE_IMAGE"]}")

    # 创建线程池
    with ThreadPoolExecutor(max_workers=setting.THREAD_POOL_SIZE) as executor:
        # 创建TCP套接字
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket:
            # 设置套接字选项，允许地址重用
            server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

            # 绑定地址和端口
            server_socket.bind((bindip, port))

            # 开始监听，设置最大连接队列
            server_socket.listen(50)

            logging.info(f"服务器已成功启动，等待客户端连接...")

            # 接受客户端连接
            while True:
                try:
                    client_socket, client_address = server_socket.accept()
                    # 提交到线程池处理
                    executor.submit(handle_client, client_socket, client_address)
                except Exception as e:
                    logging.error(f"接受客户端连接时出错: {e}", exc_info=True)


if __name__ == "__main__":
    main()